import numpy as np

BM, BN, BK, TM, TN = 128, 128, 8, 8, 8
SIMD_WIDTH = 4
assert BM * SIMD_WIDTH == BK * TM * TN
assert BM == BN
assert BK % SIMD_WIDTH == 0
assert BN % SIMD_WIDTH == 0
AV = BK // SIMD_WIDTH
BV = BN // SIMD_WIDTH
dtype = np.float32


def divup(a, b):
    return (a + b - 1) // b


def matmul_kernel(grid, block, args):
    global_a, global_b, global_c = args
    M, K = global_a.shape
    _, N = global_b.shape
    assert K % BK == 0
    gridDim_x, gridDim_y = grid
    blockDim_x, blockDim_y = block
    shared_a = np.zeros((gridDim_y, gridDim_x, BM, BK), dtype=dtype)
    shared_b = np.zeros((gridDim_y, gridDim_x, BK, BN), dtype=dtype)
    n_bk = divup(K, BK)
    global_a_ = global_a.reshape(gridDim_y, -1, n_bk, AV, SIMD_WIDTH)
    global_b_ = global_b.reshape(n_bk, -1, gridDim_x, BV, SIMD_WIDTH)

    for blockIdx_y in range(gridDim_y):
        for blockIdx_x in range(gridDim_x):
            register_c = np.zeros((blockDim_y, blockDim_x, TM, TN), dtype=dtype)
            for bk in range(0, n_bk, 1):
                for threadIdx_y in range(blockDim_y):
                    for threadIdx_x in range(blockDim_x):
                        tid = threadIdx_y * blockDim_x + threadIdx_x
                        shared_a_ = shared_a.reshape(
                            gridDim_y, gridDim_x, -1, AV, SIMD_WIDTH
                        )
                        shared_b_ = shared_b.reshape(
                            gridDim_y, gridDim_x, -1, BV, SIMD_WIDTH
                        )

                        shared_a_[
                            blockIdx_y,
                            blockIdx_x,
                            tid // AV,
                            tid % AV,
                            0:SIMD_WIDTH,
                        ] = global_a_[blockIdx_y, tid // AV, bk, tid % AV, 0:SIMD_WIDTH]
                        shared_b_[
                            blockIdx_y,
                            blockIdx_x,
                            tid // BV,
                            tid % BV,
                            0:SIMD_WIDTH,
                        ] = global_b_[bk, tid // BV, blockIdx_x, tid % BV, 0:SIMD_WIDTH]
                # syncthreads
                shared_a_ = shared_a.reshape(gridDim_y, gridDim_x, blockDim_y, TM, BK)
                shared_b_ = shared_b.reshape(gridDim_y, gridDim_x, BK, blockDim_x, TN)
                for threadIdx_y in range(blockDim_y):
                    for threadIdx_x in range(blockDim_x):
                        for k in range(BK):
                            for m in range(TM):
                                for n in range(TN):
                                    register_c[threadIdx_y, threadIdx_x, m, n] += (
                                        shared_a_[
                                            blockIdx_y, blockIdx_x, threadIdx_y, m, k
                                        ]
                                        * shared_b_[
                                            blockIdx_y, blockIdx_x, k, threadIdx_x, n
                                        ]
                                    )
                # syncthreads
            global_c_ = global_c.reshape(
                gridDim_y, blockDim_y, TM, gridDim_x, blockDim_x, TN
            )
            for threadIdx_y in range(blockDim_y):
                for threadIdx_x in range(blockDim_y):
                    for m in range(TM):
                        for n in range(0, TN, SIMD_WIDTH):
                            global_c_[
                                blockIdx_y,
                                threadIdx_y,
                                m,
                                blockIdx_x,
                                threadIdx_x,
                                n : n + SIMD_WIDTH,
                            ] = register_c[
                                threadIdx_y, threadIdx_x, m, n : n + SIMD_WIDTH
                            ]


def main():
    np.random.seed(0)
    M, N, K = 128, 256, 16
    a = np.random.uniform(-1, 1, (M, K)).astype(dtype)
    b = np.random.uniform(-1, 1, (K, N)).astype(dtype)
    c = np.zeros((M, N), dtype=dtype)
    assert N % BN == 0
    assert M % BM == 0
    grid = (divup(N, BN), divup(M, BM))
    block = (BN // TN, BM // TM)

    matmul_kernel(grid, block, (a, b, c))
    print(np.sum(np.abs(c - a @ b)))


main()
